import os
import cv2
from pathlib import Path

# 设置视频文件夹路径
video_dir = './video'
output_base_dir = './other'

# 获取所有 AVI 文件
video_files = list(Path(video_dir).glob('*.AVI'))

for idx, video_file in enumerate(video_files):
    print(f"Processing {video_file.name}")

    video_idx = int(video_file.stem)
    if video_idx < 600:
        output_base_dir = './train'
    elif video_idx < 700:
        output_base_dir = './valid'
    else:
        output_base_dir = './test'

    # 创建输出目录
    output_dir = Path(output_base_dir) / ('0' + video_file.stem) / 'images'
    output_dir.mkdir(parents=True, exist_ok=True)

    # 打开视频文件
    cap = cv2.VideoCapture(str(video_file))

    frame_count = 1
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # 保存帧为图片文件
        frame_filename = f"{frame_count:04d}.png"
        frame_path = output_dir / frame_filename
        cv2.imwrite(str(frame_path), frame)

        frame_count += 1

    cap.release()
    print(f"Extracted {frame_count} frames from {video_file.name}")